/*
 * Copyright 2016 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <iostream>
#include <string>
#include <vector>
#include <memory>

#include "rclcpp/rclcpp.hpp"
#include "rcpputils/rcpputils/filesystem_helper.hpp"
#include "rosidl_typesupport_cpp/message_type_support.hpp"
#include "rosbag2_cpp/writer.hpp"
#include "rosbag2_cpp/writers/sequential_writer.hpp"
#include "rosbag2_storage/serialized_bag_message.hpp"
#include "tf2_msgs/msg/tf_message.hpp"
#include "geometry_msgs/msg/transform_stamped.hpp"
#include "absl/strings/str_cat.h"
#include "slam/transform/transform.h"
#include "slam_ros/msg_conversion.h"
#include "slam_ros/time_conversion.h"

namespace slam_ros {

class PbstreamToBagConverter : public rclcpp::Node {
 public:
  explicit PbstreamToBagConverter(const rclcpp::NodeOptions& options)
      : Node("pbstream_to_bag_converter", options) {
    this->declare_parameter<std::string>("input", "");
    this->declare_parameter<std::string>("output", "");
    this->declare_parameter<std::string>("parent_frame", "map");

    input_filename_ = this->get_parameter("input").as_string();
    output_filename_ = this->get_parameter("output").as_string();
    parent_frame_id_ = this->get_parameter("parent_frame").as_string();

    if (input_filename_.empty() || output_filename_.empty()) {
      RCLCPP_ERROR(this->get_logger(), 
                  "Missing required parameters: input or output filename");
      rclcpp::shutdown();
      return;
    }

    ConvertPbstreamToBag();
  }

 private:
  geometry_msgs::msg::TransformStamped ToTransformStamped(
      int64_t timestamp_uts, const std::string& child_frame_id,
      const slam::transform::proto::Rigid3d& parent_T_child) {
    static int64_t seq = 0;
    geometry_msgs::msg::TransformStamped transform_stamped;
    transform_stamped.header.stamp = slam_ros::ToRos(
        ::slam::common::FromUniversal(timestamp_uts));
    transform_stamped.header.frame_id = parent_frame_id_;
    transform_stamped.child_frame_id = child_frame_id;
    transform_stamped.transform = slam_ros::ToGeometryMsgTransform(
        ::slam::transform::ToRigid3(parent_T_child));
    return transform_stamped;
  }

  void ConvertPbstreamToBag() {
    try {
        const auto pose_graph = slam::io::DeserializePoseGraphFromFile(input_filename_);
        
        // 创建rosbag2写入器
        rosbag2_cpp::Writer writer;
        rosbag2_storage::StorageOptions storage_options;
        storage_options.uri = output_filename_;
        storage_options.storage_id = "sqlite3";
        
        rosbag2_cpp::ConverterOptions converter_options;
        converter_options.input_serialization_format = "cdr";
        converter_options.output_serialization_format = "cdr";
        
        writer.open(storage_options, converter_options);

        for (const auto& trajectory : pose_graph.trajectory()) {
            const auto child_frame_id =
                absl::StrCat("trajectory_", trajectory.trajectory_id());
            
            RCLCPP_INFO(this->get_logger(),
                       "Writing tf and TransformStamped for trajectory id %d with %d nodes",
                       trajectory.trajectory_id(), trajectory.node_size());

            for (const auto& node : trajectory.node()) {
                // 创建TF消息
                auto tf_msg = std::make_shared<tf2_msgs::msg::TFMessage>();
                auto transform_stamped = ToTransformStamped(
                    node.timestamp(), child_frame_id, node.pose());
                tf_msg->transforms.push_back(transform_stamped);

                // 写入轨迹特定话题
                writer.write(
                    *tf_msg,
                    child_frame_id,
                    transform_stamped.header.stamp);

                // 写入/tf话题
                writer.write(
                    *tf_msg,
                    "/tf",
                    transform_stamped.header.stamp);
            }
        }

        RCLCPP_INFO(this->get_logger(), "Successfully converted pbstream to rosbag2");
    } catch (const std::exception& e) {
        RCLCPP_ERROR(this->get_logger(), "Conversion failed: %s", e.what());
        rclcpp::shutdown();
    }
  }

  std::string input_filename_;
  std::string output_filename_;
  std::string parent_frame_id_;
};

}  // namespace slam_ros

int main(int argc, char* argv[]) {
  rclcpp::init(argc, argv);
  rclcpp::NodeOptions options;
  auto node = std::make_shared<slam_ros::PbstreamToBagConverter>(options);
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}